-- Piotr Bober
-- Compiler Construction 2009/10
-- the stack machine

module Maszyna where

import Control.Monad.State
import Char (chr)
import Prelude hiding (div, log)
import System.IO (hFlush, stdout, hPrint, hPutStrLn, Handle)

-- stack definition
data Stos a = S Int [a]
	deriving Show

pop :: Stos a -> (a, Stos a)
pop (S 0 _) = error "pusty stos"
pop (S _ []) = error "pusty stos"
pop (S n (x:xs)) = (x, S (n-1) xs)

push :: a -> Stos a -> Stos a
push c (S n xs) = S (n+1) (c:xs)

isEmpty :: Stos a -> Bool
isEmpty (S n _) = n == 0

emptyStack :: Stos a
emptyStack = S 0 []

-- state definition

data Stan = Stan {
		pCnt :: Int,			-- program counter
		code :: Code,			-- code
		argStack :: Stos Double,	-- argument stack
		adrStack :: Stos Int		-- address stack
		}

instance Show Stan where
	show (Stan pc cs s1 s2) = "program counter: " ++ show pc
				            ++ "\ncode length: " ++ show (length cs)
				            ++ "\nargument stack: " ++ show s1
				            ++ "\naddress stack: " ++ show s2

pustyStan :: Stan
pustyStan = Stan 0 emptyCode emptyStack emptyStack

type Command = StateT Stan IO (Int,String) -- the last is the return value and the name of the command (for execution logging)

type Code = [Command]

emptyCode :: Code
emptyCode = []

---------------------------------  H O W   I T   W O R K S  --------------------------------------------------
--
-- the operation of the machine is handled by the function "run"
--
-- the machine starts operating with the program counter set to 0, the given code and empty stacks.
--
-- after execution of every command included in the code, the program counter is incremented and then control goes to the command pointed by the new value of the counter.
--
-- the return value -1 means that execution will be immediately terminated
-- the return value 0 means that a command completed successfully, and the machine will keep working
-- any other return value indicates a failure and execution will also be terminated
--
-- if the program counter ever gets bigger or equal (the code is indexed from 0 to n-1) than the length of the code then operation is also terminated
--
-- the run function passess the logging info from the functions to a given io handle
--------------------------------------------------------------------------------------------------------------

run :: Handle -> Code -> IO ()
run rh [] = hPrint rh pustyStan >> return ()
run rh c =
   let
      stzero = Stan 0 c emptyStack emptyStack
      loop stan =
         let
            pc = pCnt stan
            n = length $ code stan
         in
            if pc > n-1
               then do
                  putStrLn $ "program counter too large - execution terminated"
                  print stan
               else do
                  let cmd = code stan !! (pCnt stan)
                  ((retval,log), s) <- runStateT cmd stan
                  let stan' = Stan (pCnt s + 1) (code s) (argStack s) (adrStack s)
                  case retval of
                     -1 -> do
                              hPutStrLn rh log
                              hPutStrLn rh "\nexecution terminated successfully, the final state:"
                              hPrint rh stan'
                              return ()
                     0  -> do
                              hPutStrLn rh log
                              loop stan'
                     _  -> do
                              hPutStrLn rh log
                              hPutStrLn rh $ "\nterminated with code " ++ show retval
                              hPrint rh stan'
   in
      do
--         hPrint rh stzero
         loop stzero
--------------------------------------------------------------------------------------------------------------

-- basic commands

-- exit - returns the value -1 which terminates computation
exit :: Command
exit = StateT $ \s -> return ((-1,show (argStack s) ++ "\nexit"),s)

-- skip
skip :: Command
skip = StateT $ \s -> return ((0,show (argStack s) ++ "\nskip"),s)

-- store/load - push/pop a constant on the stack
store :: Double -> Command
store x = StateT $ \s -> return ((0,show (argStack s) ++ "\nstore " ++ show x), Stan (pCnt s) (code s) (push x $ argStack s) (adrStack s))

load :: Command
load = StateT $ \s -> return ((0,show (argStack s) ++ "\nload"), Stan (pCnt s) (code s) (snd $ pop $ argStack s) (adrStack s))

-- copy - pushes a copy of the top value of the stack
copy :: Command
copy = StateT $ \s ->
   let
      as = argStack s
      (a,_) = pop as
   in
      return ((0,show as ++ "\ncopy"), Stan (pCnt s) (code s) (push a as) (adrStack s))

-- arithmetic operations on the 2 numbers from the top of the stack, remove them, and put the result back on the stack
add, mul, sub, div :: Command

add = arithmetic "add" (+)
mul = arithmetic "mul" (*)
sub = arithmetic "sub" (-)
div = arithmetic "div" (\x y -> if y == 0.0 then 0.0 else x/y)

arithmetic :: String -> (Double -> Double -> Double) -> Command
arithmetic name f = StateT $ \s ->
   let
      as    = argStack s
      (a,b) = pop as
      (c,d) = pop b
   in
      return ((0,show as ++ '\n' : name), Stan (pCnt s) (code s) (push (f a c) d) (adrStack s))

-- negate - negate the number on top of the stack
negate :: Command
negate = StateT $ \s ->
   let
      as    = argStack s
      (a,b) = pop as
   in
      return ((0,show as ++ "\nnegate"), Stan (pCnt s) (code s) (push (-a) b) (adrStack s))

-- functions that compare the 2 numbers on top of the stack, remove them, and push the result of the comparison back on the stack
less, greater, lessequal, greaterequal, equal, notequal :: Command

less           = boolean "less" (<)
greater        = boolean "greater" (>)
lessequal      = boolean "lessequal" (<=)
greaterequal   = boolean "greaterequal" (>=)
equal          = boolean "equal" (==)
notequal       = boolean "notequal" (/=)

boolean :: String -> (Double -> Double -> Bool) -> Command
boolean name f = StateT $ \s ->
   let
      as    = argStack s
      (a,b) = pop as
      (c,d) = pop b
   in if f a c
         then return ((0,show as ++ '\n' : name), Stan (pCnt s) (code s) (push 1.0 d) (adrStack s))
         else return ((0,show as ++ '\n' : name), Stan (pCnt s) (code s) (push 0.0 d) (adrStack s))

-- read, write - get the n-th value from the stack and push it on top / get the value on top of the stack and put it on the n-th position
read', write :: Int -> Command

read' n | n < 0 = StateT $ \s -> putStrLn "read : negative index" >> return ((1,show (argStack s) ++ "\nread " ++ show n ++ " - error"),s)
read' n  = StateT $ \s ->
   let
      as@(S len xs) = argStack s
      (a,x,b) = del n xs
   in
      if len < n
         then putStrLn "read: index out of range" >> return ((1,show as ++ "\nread " ++ show n ++ " - error"),s)
         else return ((0,show as ++ "\nread " ++ show n),Stan (pCnt s) (code s) (S len (x ++ a ++ b)) (adrStack s))

write n | n < 0 = StateT $ \s -> putStrLn "write: negative index" >> return ((1,show (argStack s) ++ "\nwrite " ++ show n ++ " - error"),s)
write n = StateT $ \s ->
   let
      as@(S len xs) = argStack s
      (a,b) = splitAt (n-1) $ tail xs
   in
      if len < n
         then putStrLn "write: index out of range" >> return ((1,show as ++ "\nwrite " ++ show n ++ " - error"),s)
         else return ((0,show as ++ "\nwrite " ++ show n), Stan (pCnt s) (code s) (S len (a ++ [head xs] ++  b)) (adrStack s))

-- jumps and calls
jump, jumpnonzero, jumpzero, call :: Int -> Command

jump n = StateT $ \s ->
   let
      as = argStack s
      pc = pCnt s
      cd = code s
      cl = length cd
      newpc = pc + n - 1
   in
      if newpc >= cl || newpc < 0
         then putStrLn "jump: address out of range" >> return ((1,show as ++ "\njump " ++ show n ++ " - error"),s)
         else return ((0,show as ++ "\njump " ++ show n), Stan newpc cd as (adrStack s))

jumpnonzero n = StateT $ \s ->
   let
      as = argStack s
      pc = pCnt s
      cd = code s
      cl = length cd
      newpc = pc + n - 1
   in
      if newpc >= cl || newpc < 0
         then putStrLn "jumpnonzero: address out of range" >> return ((1,show as ++ "\njumpnonzero - error"),s)
         else
            let
               (a,_) = pop as
            in
               if a /= 0.0
                  then return ((0,show as ++ "\njumpnonzero " ++ show n ++ " (jump)"), Stan newpc cd as (adrStack s))
                  else return ((0,show as ++ "\njumpnonzero " ++ show n ++ " (no jump)"), s)

jumpzero n = StateT $ \s ->
   let
      as    = argStack s
      pc = pCnt s
      cd = code s
      cl = length cd
      newpc = pc + n - 1
   in
      if newpc >= cl || newpc < 0
         then putStrLn "jumpzero: address out of range" >> return ((1,show as ++ "\njumpzero - error"),s)
         else
            let
               (a,_) = pop as
            in
               if a == 0.0
                  then return ((0,show as ++ "\njumpzero " ++ show n ++ " (jump)"), Stan newpc cd as (adrStack s))
                  else return ((0,show as ++ "\njumpzero " ++ show n ++ " (no jump)"), s)

call n = StateT $ \s ->
   let
      as = argStack s
      pc = pCnt s
      cd = code s
      cl = length cd
      newpc = pc + n - 1
   in
      if newpc >= cl || newpc < 0
         then putStrLn ("call: address " ++ show newpc ++ " out of range") >> return ((1,show as ++ "\ncall " ++ show n ++ " - error"),s)
         else return ((0,show as ++ "\ncall " ++ show n ++ " (return to " ++ show pc ++ ")"), Stan newpc cd as (push pc $ adrStack s))

ret :: Command
ret = StateT $ \s ->
   let
      as    = argStack s
      adrs  = adrStack s
      (a,b) = pop adrs
   in return ((0,show as ++ "\nret (to " ++ show a ++ ")"), Stan (a) (code s) (as) (b))

-- input/output
input, output, outputchar :: Command

input = StateT $ \s -> do
                           hFlush stdout
                           line <- getLine
                           catch
                              (do
                                 let x = read line :: Double
                                 let as = argStack s
                                 return ((0,show as ++ "\ninput"), Stan (pCnt s) (code s) (push x as) (adrStack s)))
                              (const $ putStrLn "input: no parse of input" >> return ((1,show (argStack s) ++ "\ninput - error"),s))

output = StateT $ \s ->
   let
      as    = argStack s
      (a,_) = pop as
   in
      do
         print a
         return ((0,show as ++ "\noutput"), Stan (pCnt s) (code s) as (adrStack s))

outputchar = StateT $ \s ->
   let
      as    = argStack s
      (a,_) = pop as
   in
      catch
         (do
            putChar $ chr $ floor a
            return ((0,show as ++ "\noutputchar"), Stan (pCnt s) (code s) as (adrStack s)))
         (const $ putStrLn "outputchar: error converting number to character" >> return ((1,show as ++ "\noutputchar"),s))

-- auxiliary functions
del :: Int -> [a] -> ([a],[a],[a])
del n s | n > length s = error "del: index out of range"
del n s =
   let
      foo ak _ [] = (reverse ak,[],[])
      foo ak 1 (x:xs) = (reverse ak,[x],xs)
      foo ak k (x:xs) = foo (x:ak) (k-1) xs
   in
      foo [] n s

-- testing
test, test1, test2, test3, test4, test5 :: IO ()

test = run stdout $ map store [1..10] ++ [sub, write 4, read' 5, exit]

test1 = run stdout $ map store (reverse [1..10]) ++ take 9 (repeat add) ++ [exit]

test2 = run stdout $ map store (reverse [1..5]) ++ [mul, mul, mul, mul, exit]

test3 = run stdout [store 25, store 625, div, exit]

test4 = run stdout [
   store 5,
   store 1,
   store 0,
   jumpzero 3,
   add,
   store 10,
   store (-1),
   exit]

test5 = run stdout [
   store 5,
   store 1,
   store 0,
   jumpnonzero 3,
   add,
   store 10,
   store (-1),
   exit]

-- silnia
silnia :: Int -> IO ()
{-
input n
i = n;
f = 1;
while (i>0) {
   f = f * i
   i = i - 1
}
output f
-}
silnia n = run stdout [
   store 1,    -- f = 1
   store n',   -- i <- input n
   jumpzero 8, -- while i > 0
   copy,       -- temp = i
   read' 3,    -- move f to the top
   mul,        -- f = f * temp (only f stays at the stack)
   write 2,    -- move f back to the bottom
   store (-1),
   add,        -- i = i - 1
   jump (-7),  -- end of while
   load,
   exit]
   where n' = fromIntegral n

-- interactive - unfinished
{-
main :: IO ()
main = loop pustyStan

loop :: Stan -> IO ()
loop stan =
   do
      putStr "> "
      line <- getLine
      let cmdargs = words line
      if cmdargs == []
         then loop stan
         else
            let
               cmd = head cmdargs
            in case cmd of
               "quit" -> return ()
               "init" -> loop pustyStan
               "show" -> do
                  print stan
                  loop stan
               "push" -> do
                  let arg = cmdargs !! 1
                  (retval, stan') <- runStateT (store (read arg :: Double)) stan
                  print retval
                  loop stan'
               _      -> do
                  putStrLn $ "unknown command: " ++ cmd
                  loop stan
-}
